#include <math.h>
#include <stdio.h>

#define EPSILON 1.0e-3

void check_rotary_embedding(int64_t *pos, float *query, int64_t base,
		int64_t rot_dim, int64_t num_tokens,
		int64_t num_heads, int64_t head_size) {
	for (int token_id = 0; token_id < num_tokens; ++token_id) {
		int64_t token_pos = pos[token_id];
		for (int head_id = 0; head_id < num_heads; ++head_id) {
			int64_t token_head_pos =
				token_id * num_heads * head_size + head_id * head_size;
			for (int rot_id = 0; rot_id < rot_dim / 2; ++rot_id) {
				int64_t x_index = rot_id;
				int64_t y_index = rot_dim / 2 + rot_id;
				float x = query[token_head_pos + x_index];
				float y = query[token_head_pos + y_index];
				float cos_val = cos(
						token_pos * (1 / pow((float)base, 2 * rot_id / (float)rot_dim)));
				float sin_val = sin(
						token_pos * (1 / pow((float)base, 2 * rot_id / (float)rot_dim)));
				query[token_head_pos + x_index] = x * cos_val - y * sin_val;
				query[token_head_pos + y_index] = y * cos_val + x * sin_val;
			}
		}
	}
}

__global__ void rotary_embedding_kernel(int64_t *pos, float *query,
		int64_t base, int64_t rot_dim,
		int64_t num_tokens, int64_t num_heads,
		int64_t head_size) {
	int token_id = blockIdx.x;
	int64_t token_pos = pos[token_id];
	for (int head_id = threadIdx.x; head_id < num_heads; head_id += blockDim.x) {
		int64_t token_head_pos =
			token_id * num_heads * head_size + head_id * head_size;
		for (int rot_id = 0; rot_id < rot_dim / 2; ++rot_id) {
			int64_t x_index = rot_id;
			int64_t y_index = rot_dim / 2 + rot_id;
			float x = query[token_head_pos + x_index];
			float y = query[token_head_pos + y_index];
			float cos_val =
				cos(token_pos * (1 / pow((float)base, 2 * rot_id / (float)rot_dim)));
			float sin_val =
				sin(token_pos * (1 / pow((float)base, 2 * rot_id / (float)rot_dim)));
			query[token_head_pos + x_index] = x * cos_val - y * sin_val;
			query[token_head_pos + y_index] = y * cos_val + x * sin_val;
		}
	}
}


// pos: pos, [num_tokens]
// query: query, [num_tokens, num_heads * head_size]
// base: \theta = base^(-2*i/rot_dim)
// rot_dim: \theta = base^(-2*i/rot_dim)
// num_tokens: The num of tokens.
// num_heads: The num of attention heads.
// head_size: The size of each attention head.
void rotary_embedding(int64_t *pos, float *query, int64_t base, int64_t rot_dim,
		int64_t num_tokens, int64_t num_heads,
		int64_t head_size) {
	dim3 grid(num_tokens);
	dim3 block(num_heads);
	rotary_embedding_kernel<<<grid, block>>>(pos, query, base, rot_dim,
			num_tokens, num_heads, head_size);
}

int main() {
	const int64_t num_tokens = 128;
	const int64_t num_heads = 16;
	const int64_t head_size = 128;
	const int64_t base = 10000;
	const int64_t rot_dim = 128;
	const int64_t max_position_embeddings = 327;

	int64_t *h_pos = (int64_t *)malloc(num_tokens * sizeof(int64_t));
	float *h_query =
		(float *)malloc(num_tokens * num_heads * head_size * sizeof(float));
	float *check_query =
		(float *)malloc(num_tokens * num_heads * head_size * sizeof(float));

	for (int n = 0; n < num_tokens; ++n) {
		h_pos[n] = rand() % (max_position_embeddings + 1);
	}
	for (int n = 0; n < num_tokens * num_heads * head_size; ++n) {
		h_query[n] = rand() / (float)RAND_MAX;
		check_query[n] = h_query[n];
	}

	int64_t *d_pos;
	float *d_query;
	cudaMalloc((void **)&d_pos, num_tokens * sizeof(int64_t));
	cudaMalloc((void **)&d_query,
			num_tokens * num_heads * head_size * sizeof(float));

	cudaMemcpy(d_pos, h_pos, num_tokens * sizeof(int64_t),
			cudaMemcpyHostToDevice);
	cudaMemcpy(d_query, h_query,
			num_tokens * num_heads * head_size * sizeof(float),
			cudaMemcpyHostToDevice);

	cudaEvent_t start, stop;
	cudaEventCreate(&start);
	cudaEventCreate(&stop);

	cudaEventRecord(start);

	rotary_embedding(d_pos, d_query, base, rot_dim, num_tokens, num_heads,
			head_size);

	cudaEventRecord(stop);
	cudaEventSynchronize(stop);

	float time = 0;
	cudaEventElapsedTime(&time, start, stop);
	printf("time: %f ms\n", time);

	cudaMemcpy(h_query, d_query,
			num_tokens * num_heads * head_size * sizeof(float),
			cudaMemcpyDeviceToHost);

	check_rotary_embedding(h_pos, check_query, base, rot_dim, num_tokens,
			num_heads, head_size);
	int has_error = 0;
	for (int n = 0; n < num_tokens * num_heads * head_size; ++n) {
		if (fabs(h_query[n] - check_query[n]) > EPSILON) {
			has_error = 1;
		}
	}
	printf("%s\n", has_error ? "Has errors" : "No errors");

	free(h_pos);
	free(h_query);
	free(check_query);
	cudaFree(d_pos);
	cudaFree(d_query);

	return 0;
}
